from __future__ import annotations
from typing_extensions import Literal
from typing import Any, Callable, Dict, List, Optional, Tuple, Union

TickReason = Literal["Unknown"]

def get_version() -> str:
    """Get the version of the BAML Python client."""
    ...

def set_log_level(
    level: Literal["DEBUG", "INFO", "WARN", "ERROR", "OFF"] | str,
) -> None:
    """Set the log level for the BAML Python client."""
    ...

def set_log_json_mode(use_json: bool) -> None:
    """Set the log JSON mode for the BAML Python client."""
    ...

def get_log_level() -> str:
    """Get the log level for the BAML Python client."""
    ...

def set_log_max_chunk_length(length: int) -> None:
    """Set the maximum log chunk length for the BAML Python client."""
    ...

class AbortController:
    """Controller for cancelling BAML operations."""

    def __init__(self, timeout_ms: Optional[int] = None) -> None:
        """
        Creates a new abort controller with an optional timeout in milliseconds.
        Once aborted, the AbortController will forever remain in an an aborted state.
        The timeout will only start AFTER the object is passed to a BAML function.

        Args:
            timeout_ms: The timeout in milliseconds. If not provided, AbortController will not timeout.
        """
        ...

    def abort(self) -> None:
        """
        Immediately abort all operations.
        """
        ...

    @property
    def aborted(self) -> bool:
        """
        Check the state of this controller.
        Once aborted, the AbortController will forever remain in an an aborted state.
        """
        ...

class FunctionResult:
    """The result of a BAML function call.

    Represents any of:

        - a successful LLM call, with a successful type parse
        - a successful LLM call, with a failed type parse
        - a failed LLM call, due to a provider outage or other network error
        - a failed LLM call, due to an inability to build the request
        - or any other outcome, really

    We only expose the parsed result to Python right now.
    """

    def __str__(self) -> str: ...
    # Returns True if the function call was successful, False otherwise
    def is_ok(self) -> bool: ...
    def cast_to(
        self,
        enum_module: Any,
        class_module: Any,
        partial_class_module: Any,
        allow_partials: bool,
        runtime: BamlRuntime,
    ) -> Any: ...

    # This is a debug function that returns the internal representation of the response
    # This is not to be relied upon and is subject to change
    # Usage:
    #   result = await runtime.call_function(...)
    #   val = json.loads(result.unstable_internal_repr())
    #   print(val)
    def unstable_internal_repr(self) -> str: ...

class FunctionResultStream:
    """The result of a BAML function stream.

    Provides a callback interface to receive events from a BAML result stream.

    Use `on_event` to set the callback, and `done` to drive the stream to completion.
    """

    def __str__(self) -> str: ...
    def on_event(
        self, on_event: Callable[[FunctionResult], None]
    ) -> FunctionResultStream: ...
    async def done(self, ctx: RuntimeContextManager) -> FunctionResult: ...

class SyncFunctionResultStream:
    """The result of a BAML function stream.

    Provides a callback interface to receive events from a BAML result stream.

    Use `on_event` to set the callback, and `done` to drive the stream to completion.
    """

    def __str__(self) -> str: ...
    def on_event(
        self, on_event: Callable[[FunctionResult], None]
    ) -> SyncFunctionResultStream: ...
    def done(self, ctx: RuntimeContextManager) -> FunctionResult: ...

class BamlImagePy:
    @staticmethod
    def from_url(url: str) -> BamlImagePy: ...
    @staticmethod
    def from_base64(media_type: str, base64: str) -> BamlImagePy: ...
    def is_url(self) -> bool: ...
    def is_base64(self) -> bool: ...
    def as_url(self) -> str: ...
    def as_base64(self) -> Tuple[str, str]: ...

class BamlAudioPy:
    @staticmethod
    def from_url(url: str) -> BamlAudioPy: ...
    @staticmethod
    def from_base64(media_type: str, base64: str) -> BamlAudioPy: ...
    def is_url(self) -> bool: ...
    def is_base64(self) -> bool: ...
    def as_url(self) -> str: ...
    def as_base64(self) -> Tuple[str, str]: ...

class BamlPdfPy:
    @staticmethod
    def from_url(url: str) -> BamlPdfPy: ...
    @staticmethod
    def from_base64(base64: str) -> BamlPdfPy: ...
    def is_url(self) -> bool: ...
    def is_base64(self) -> bool: ...
    def as_url(self) -> str: ...
    def as_base64(self) -> Tuple[str, str]: ...

class BamlVideoPy:
    @staticmethod
    def from_url(url: str) -> BamlVideoPy: ...
    @staticmethod
    def from_base64(media_type: str, base64: str) -> BamlVideoPy: ...
    def is_url(self) -> bool: ...
    def is_base64(self) -> bool: ...
    def as_url(self) -> str: ...
    def as_base64(self) -> Tuple[str, str]: ...

class RuntimeContextManager:
    def upsert_tags(self, tags: Dict[str, Any]) -> None: ...
    def deep_clone(self) -> RuntimeContextManager: ...
    def context_depth(self) -> int: ...

class BamlRuntime:
    def __getstate__(self) -> Tuple[str, Dict[str, str]]: ...
    def __setstate__(self, state: Tuple[str, Dict[str, str]]) -> None: ...
    @staticmethod
    def from_directory(directory: str, env_vars: Dict[str, str]) -> BamlRuntime: ...
    async def call_function(
        self,
        function_name: str,
        args: Dict[str, Any],
        ctx: RuntimeContextManager,
        tb: Optional[TypeBuilder],
        cr: Optional[ClientRegistry],
        collectors: List[Collector],
        env_vars: Dict[str, str],
        tags: Optional[Dict[str, str]] = None,
        abort_controller: Optional[AbortController] = None,
    ) -> FunctionResult: ...
    def call_function_sync(
        self,
        function_name: str,
        args: Dict[str, Any],
        ctx: RuntimeContextManager,
        tb: Optional[TypeBuilder],
        cr: Optional[ClientRegistry],
        collectors: List[Collector],
        env_vars: Dict[str, str],
        tags: Optional[Dict[str, str]] = None,
        abort_controller: Optional[AbortController] = None,
    ) -> FunctionResult: ...
    @staticmethod
    def from_files(
        root_path: str, files: Dict[str, str], env_vars: Dict[str, str]
    ) -> BamlRuntime: ...
    def reset(
        self, root_path: str, files: Dict[str, str], env_vars: Dict[str, str]
    ) -> None: ...
    def stream_function(
        self,
        function_name: str,
        args: Dict[str, Any],
        on_event: Optional[Callable[[FunctionResult], None]],
        ctx: RuntimeContextManager,
        tb: Optional[TypeBuilder],
        cr: Optional[ClientRegistry],
        collectors: List[Collector],
        env_vars: Dict[str, str],
        tags: Optional[Dict[str, str]] = None,
        on_tick: Optional[Callable[[], None]] = None,
    ) -> FunctionResultStream: ...
    def stream_function_sync(
        self,
        function_name: str,
        args: Dict[str, Any],
        on_event: Optional[Callable[[FunctionResult], None]],
        ctx: RuntimeContextManager,
        tb: Optional[TypeBuilder],
        cr: Optional[ClientRegistry],
        collectors: List[Collector],
        env_vars: Dict[str, str],
        tags: Optional[Dict[str, str]] = None,
        on_tick: Optional[Callable[[], None]] = None,
    ) -> SyncFunctionResultStream: ...
    def create_context_manager(self) -> RuntimeContextManager: ...
    def flush(self) -> None: ...
    def drain_stats(self) -> TraceStats: ...
    def set_log_event_callback(
        self, handler: Optional[Callable[[BamlLogEvent], None]]
    ) -> None: ...
    async def build_request(
        self,
        function_name: str,
        args: Dict[str, Any],
        ctx: RuntimeContextManager,
        tb: Optional[TypeBuilder],
        cr: Optional[ClientRegistry],
        env_vars: Dict[str, str],
        is_stream: bool,
    ) -> HTTPRequest: ...
    def build_request_sync(
        self,
        function_name: str,
        args: Dict[str, Any],
        ctx: RuntimeContextManager,
        tb: Optional[TypeBuilder],
        cr: Optional[ClientRegistry],
        env_vars: Dict[str, str],
        is_stream: bool,
    ) -> HTTPRequest: ...
    def parse_llm_response(
        self,
        function_name: str,
        llm_response: str,
        enum_module: Any,
        cls_module: Any,
        partial_cls_module: Any,
        allow_partials: bool,
        ctx: RuntimeContextManager,
        tb: Optional[TypeBuilder],
        cr: Optional[ClientRegistry],
        env_vars: Dict[str, str],
    ) -> Any: ...
    def disassemble(self, function_name: str) -> None: ...

class LogEventMetadata:
    event_id: str
    parent_id: Optional[str]
    root_event_id: str

    def __init__(
        self, event_id: str, parent_id: Optional[str], root_event_id: str
    ) -> None: ...

class BamlLogEvent:
    metadata: LogEventMetadata
    prompt: Optional[str]
    raw_output: Optional[str]
    parsed_output: Optional[str]
    start_time: str

    def __init__(
        self,
        metadata: LogEventMetadata,
        prompt: Optional[str],
        raw_output: Optional[str],
        parsed_output: Optional[str],
        start_time: str,
    ) -> None: ...

class TraceStats:
    @property
    def failed(self) -> int: ...
    @property
    def started(self) -> int: ...
    @property
    def finalized(self) -> int: ...
    @property
    def submitted(self) -> int: ...
    @property
    def sent(self) -> int: ...
    @property
    def done(self) -> int: ...

class BamlSpan:
    @staticmethod
    def new(
        runtime: BamlRuntime,
        function_name: str,
        args: Dict[str, Any],
        ctx: RuntimeContextManager,
    ) -> BamlSpan: ...
    def finish(self, result: Any, ctx: RuntimeContextManager) -> str | None: ...

class TypeBuilder:
    def __init__(self) -> None: ...
    def reset(self) -> None: ...
    def enum(self, name: str) -> EnumBuilder: ...
    def class_(self, name: str) -> ClassBuilder: ...
    def string(self) -> FieldType: ...
    def literal_string(self, value: str) -> FieldType: ...
    def literal_int(self, value: int) -> FieldType: ...
    def literal_bool(self, value: bool) -> FieldType: ...
    def int(self) -> FieldType: ...
    def float(self) -> FieldType: ...
    def bool(self) -> FieldType: ...
    def list(self, element_type: FieldType) -> FieldType: ...
    def null(self) -> FieldType: ...
    def optional(self, inner_type: FieldType) -> FieldType: ...
    def map(self, key_type: FieldType, value_type: FieldType) -> FieldType: ...
    def union(self, *types: FieldType) -> FieldType: ...
    def add_baml(self, baml: str, rt: BamlRuntime) -> None: ...

class Collector:
    def __init__(self, name: Optional[str] = None) -> None: ...
    @property
    def logs(self) -> List[FunctionLog]: ...
    @property
    def last(self) -> Optional[FunctionLog]: ...
    @property
    def usage(self) -> Usage: ...
    def id(self, function_log_id: str) -> Optional[FunctionLog]: ...
    def clear(self) -> None: ...
    # For debugging
    @staticmethod
    def __function_call_count() -> int: ...
    @staticmethod
    def __print_storage() -> None: ...

class FunctionLog:
    def __init__(
        self, name: str, args: Dict[str, Any], result: FunctionResult
    ) -> None: ...
    def __str__(self) -> str: ...
    @property
    def id(self) -> str: ...
    @property
    def function_name(self) -> str: ...
    @property
    def log_type(self) -> Literal["call", "stream"]: ...
    @property
    def timing(self) -> Timing: ...
    @property
    def usage(self) -> Usage: ...
    @property
    def calls(self) -> List[Union[LLMCall, LLMStreamCall]]: ...
    @property
    def raw_llm_response(self) -> Optional[str]: ...
    @property
    def metadata(self) -> Dict[str, Any]: ...
    @property
    def tags(self) -> Dict[str, Any]: ...
    @property
    def selected_call(self) -> Optional[Union[LLMCall, LLMStreamCall]]: ...

class LLMCall:
    def __init__(
        self,
        client_name: str,
        provider: str,
        timing: Timing,
        request: Dict[str, Any],
        response: Dict[str, Any],
        usage: Usage,
    ) -> None: ...
    @property
    def selected(self) -> bool: ...
    @property
    def http_request(self) -> Optional[HTTPRequest]: ...
    @property
    def http_response(self) -> Optional[HTTPResponse]: ...
    @property
    def usage(self) -> Usage: ...
    @property
    def timing(self) -> Timing: ...
    @property
    def provider(self) -> str: ...
    # The baml client name / clientregistry name
    @property
    def client_name(self) -> str: ...

class LLMStreamCall:
    def __init__(
        self,
        client_name: str,
        provider: str,
        timing: StreamTiming,
        request: Dict[str, Any],
        response: Dict[str, Any],
        usage: Usage,
    ) -> None: ...
    @property
    def selected(self) -> bool: ...
    @property
    def http_request(self) -> Optional[HTTPRequest]: ...
    @property
    def http_response(self) -> Optional[HTTPResponse]: ...
    @property
    def usage(self) -> Usage: ...

    @property
    def provider(self) -> str: ...

    # The baml client name / clientregistry name
    @property
    def client_name(self) -> str: ...
    @property
    def timing(self) -> StreamTiming: ...
    def sse_responses(self) -> Optional[List[SSEResponse]]: ...

class SSEResponse:
    @property
    def text(self) -> str: ...
    def json(self) -> Optional[Any]: ...

class Usage:
    def __init__(
        self, input_tokens: Optional[int], output_tokens: Optional[int]
    ) -> None: ...
    @property
    def input_tokens(self) -> Optional[int]: ...
    @property
    def output_tokens(self) -> Optional[int]: ...
    @property
    def cached_input_tokens(self) -> Optional[int]: ...

class Timing:
    def __init__(
        self,
        start_time: str,
        end_time: str,
        duration: Optional[int],
        time_to_first_parsed_ms: Optional[int],
    ) -> None: ...
    @property
    def start_time_utc_ms(self) -> int: ...
    @property
    def duration_ms(self) -> Optional[int]: ...

class StreamTiming(Timing):
    def __init__(
        self,
        start_time_utc_ms: int,
        duration_ms: int,
        time_to_first_parsed_ms: int,
        time_to_first_token_ms: int,
    ) -> None: ...
    @property
    def start_time_utc_ms(self) -> int: ...
    @property
    def duration_ms(self) -> Optional[int]: ...

class HTTPRequest:
    def __init__(
        self, url: str, method: str, headers: Dict[str, Any], body: str
    ) -> None: ...
    @property
    def id(self) -> str: ...
    @property
    def url(self) -> str: ...
    @property
    def method(self) -> str: ...
    @property
    def headers(self) -> Dict[str, Any]: ...
    @property
    def body(self) -> HTTPBody: ...

class HTTPBody:
    def __init__(self, raw: bytes) -> None: ...
    def raw(self) -> bytes: ...
    def text(self) -> str: ...
    def json(self) -> Dict[str, Any]: ...

class HTTPResponse:
    def __init__(self, status: int, headers: Dict[str, Any], body: str) -> None: ...
    @property
    def status(self) -> int: ...
    @property
    def headers(self) -> Dict[str, Any]: ...
    @property
    def body(self) -> HTTPBody: ...

class ClientRegistry:
    def __init__(self) -> None: ...
    def add_llm_client(
        self,
        name: str,
        provider: str,
        options: Dict[str, Any],
        retry_policy: Optional[str] = None,
    ) -> None: ...
    def set_primary(self, name: str) -> None: ...

class FieldType:
    def list(self) -> FieldType: ...
    def optional(self) -> FieldType: ...
    def __eq__(self, other: object) -> bool: ...

class EnumBuilder:
    def value(self, name: str) -> EnumValueBuilder: ...
    def alias(self, alias: Optional[str]) -> EnumBuilder: ...
    def field(self) -> FieldType: ...

class EnumValueBuilder:
    def alias(self, alias: Optional[str]) -> EnumValueBuilder: ...
    def skip(self, skip: Optional[bool] = True) -> EnumValueBuilder: ...
    def description(self, description: Optional[str]) -> EnumValueBuilder: ...

class ClassBuilder:
    def field(self) -> FieldType: ...
    def list_properties(self) -> List[Tuple[str, ClassPropertyBuilder]]: ...
    def property(self, name: str) -> ClassPropertyBuilder: ...
    def remove_property(self, name: str) -> None: ...
    def reset(self) -> None: ...

class ClassPropertyBuilder:
    def type(self, field_type: FieldType) -> ClassPropertyBuilder: ...
    def get_type(self) -> FieldType: ...
    def alias(self, alias: Optional[str]) -> ClassPropertyBuilder: ...
    def description(self, description: Optional[str]) -> ClassPropertyBuilder: ...

def invoke_runtime_cli() -> int: ...

class BamlError(Exception):
    """Base class for all BAML-related errors."""

    ...

class BamlInvalidArgumentError(BamlError):
    """Raised when an invalid argument is provided to a function."""

    ...

class BamlClientError(BamlError):
    """Raised for general client errors."""

    ...

class BamlClientHttpError(BamlClientError):
    """Raised for HTTP-related client errors."""

    ...

class BamlAbortError(BamlError):
    """Raised when a BAML operation is cancelled (abort)."""

    ...
